package v1

import (
	"context"
	"encoding/json"
	"log/slog"
	"strings"
	"time"

	"quipus/internal/database/model"
	llm "quipus/internal/llm"
	"quipus/internal/llm/factory"
	"quipus/internal/rag/rag"
	"quipus/internal/server/process"
	"quipus/internal/utils"
	v1pb "quipus/proto/gen/api/v1"
	apb "quipus/proto/gen/assist"
	cpb "quipus/proto/gen/common"

	"github.com/gofiber/websocket/v2"
	"github.com/google/cel-go/cel"
	"github.com/lithammer/shortuuid/v4"
	"github.com/pkg/errors"
	expr "google.golang.org/genproto/googleapis/api/expr/v1alpha1"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/emptypb"
)

var (
	ragBuildMap = map[string]*rag.RagBuilder{}
)

func init() {
}

func (s *APIV1Service) CreateGranary(ctx context.Context, request *v1pb.CreateGranaryRequest) (*v1pb.Granary, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	granary := convertGranaryToDB(request)
	granary.TenantID = currentUser.TenantID
	granary.CreatorID = currentUser.ID
	grs := []*model.GranaryResource{}
	if request.Sources != nil || len(request.Sources) > 1 {
		rs := []int32{}

		for _, source := range request.Sources {
			rid, e := utils.ExtractIDFromUrn(utils.ResourceUrnPrefix, source)
			if e != nil {
				return nil, e
			}
			gr := &model.GranaryResource{
				ResourceID: rid,
				Status:     v1pb.GranaryStatus_CREATED.String(),
			}
			rs = append(rs, rid)
			grs = append(grs, gr)
		}
		cb, e := json.Marshal(rs)
		if e != nil {
			return nil, e
		}
		granary.Sources = string(cb)
	}

	err = s.Gds.Handler.CreateGranary(ctx, granary)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to create granay, error: %+v", err)
	}

	for _, gr := range grs {
		gr.GranaryID = granary.ID
		err = s.Gds.Handler.UpsertGranaryResource(ctx, gr)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to create granay resource, error: %+v", err)
		}
	}

	return convertGranaryFromDB(granary), nil
}

func convertGranaryToDB(req *v1pb.CreateGranaryRequest) *model.Granary {
	temp := &model.Granary{
		Name:        req.Name,
		AvatarURL:   req.Avatar,
		Description: req.Description,
		Status:      v1pb.GranaryStatus_CREATED.String(),
		DocID:       shortuuid.New(),
	}

	if req.Tags != nil || len(req.Tags) > 1 {
		cb, e := json.Marshal(req.Tags)
		if e != nil {
			return nil
		}
		temp.Tags = string(cb)
	}

	if req.Options != nil || len(req.Options) > 1 {
		cb, e := json.Marshal(req.Options)
		if e != nil {
			return nil
		}
		temp.Options = string(cb)
	}

	if req.LlmModel != "" {
		llmId, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, req.LlmModel)
		if err != nil {
			return nil
		}
		temp.LlmModelID = llmId
	}

	if req.EmbeddingModel != "" {
		embeddingId, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, req.EmbeddingModel)
		if err != nil {
			return nil
		}
		temp.EmbeddingModelID = embeddingId
	}

	return temp
}

func convertGranaryFromDB(granary *model.Granary) *v1pb.Granary {
	temp := &v1pb.Granary{
		Urn:         utils.BuildUrn(utils.GranaryUrnPrefix, granary.ID),
		Name:        granary.Name,
		Description: granary.Description,
		Avatar:      granary.AvatarURL,
		Summary:     granary.Summary,
		Status:      v1pb.GranaryStatus(v1pb.GranaryStatus_value[granary.Status]),
		Tenant:      utils.BuildUrn(utils.TenantUrnPrefix, granary.TenantID),
	}

	if granary.Tags != "{}" {
		var tags []string
		e := json.Unmarshal([]byte(granary.Tags), &tags)
		if e != nil {
			return nil
		}
		temp.Tags = tags
	}

	if granary.Sources != "" && granary.Sources != "{}" {
		var rs []int32
		var sources []string
		e := json.Unmarshal([]byte(granary.Sources), &rs)
		if e != nil {
			return nil
		}
		for _, rid := range rs {
			sources = append(sources, utils.BuildUrn(utils.ResourceUrnPrefix, rid))
		}
		temp.Sources = sources
	}

	if granary.Options != "{}" {
		var options map[string]string
		// var options map[string]interface{}
		e := json.Unmarshal([]byte(granary.Options), &options)
		if e != nil {
			return nil
		}
		temp.Options = options
	}

	if granary.LlmModelID > 0 {
		temp.LlmModel = utils.BuildUrn(utils.ModelUrnPrefix, granary.LlmModelID)
	}

	if granary.EmbeddingModelID > 0 {
		temp.EmbeddingModel = utils.BuildUrn(utils.ModelUrnPrefix, granary.EmbeddingModelID)
	}

	return temp
}

func (s *APIV1Service) updateGranaryResources(ctx context.Context, granaryId int32, sources []int32) error {
	if sources == nil || len(sources) == 0 {
		err := s.Gds.Handler.DeleteGranaryResource(ctx, &apb.FindGranaryResource{Gid: granaryId})
		if err != nil {
			return err
		}
	}

	grs, err := s.Gds.Handler.ListGranaryResources(ctx, &apb.FindGranaryResource{Gid: granaryId})
	if err != nil {
		return err
	}

	drs := []int32{}
	js := []int32{}
	for _, gr := range grs {
		if !utils.Contains(sources, gr.ResourceID) {
			if !utils.Contains(js, gr.ResourceID) {
				drs = append(drs, gr.ResourceID)
			}
		} else {
			js = append(js, gr.ResourceID)
			sources = utils.RemoveElement(sources, gr.ResourceID)
		}
	}
	for _, gr := range drs {
		err = s.Gds.Handler.DeleteGranaryResource(ctx, &apb.FindGranaryResource{Gid: granaryId, Rid: gr})
		if err != nil {
			return err
		}
	}

	for _, rid := range sources {
		gr := &model.GranaryResource{
			GranaryID:  granaryId,
			ResourceID: rid,
			Status:     v1pb.GranaryStatus_CREATED.String(),
		}

		s.Gds.Handler.UpsertGranaryResource(ctx, gr)
	}

	return nil
}

func (s *APIV1Service) UpdateGranary(ctx context.Context, request *v1pb.UpdateGranaryRequest) (*v1pb.Granary, error) {

	req := request.Granary
	granaryId, err := utils.ExtractIDFromUrn(utils.GranaryUrnPrefix, req.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid user name: %v", err)
	}
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	if request.UpdateMask == nil || len(request.UpdateMask.Paths) == 0 {
		return nil, status.Errorf(codes.InvalidArgument, "update mask is empty")
	}

	granary, err := s.Gds.Handler.GetGranary(ctx, granaryId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}
	if granary == nil {
		return nil, status.Errorf(codes.NotFound, "user not found")
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() {
		if granary.TenantID != currentUser.TenantID {
			return nil, status.Errorf(codes.PermissionDenied, "permission denied")
		}

		if granary.CreatorID == currentUser.ID || currentUser.Role == cpb.Role_TENANT_ADMIN.String() {
			slog.Info("update granary", "granary", granary.Name)
		} else {
			return nil, status.Errorf(codes.PermissionDenied, "permission denied")
		}
	}

	currentTs := time.Now().Unix()
	update := &apb.UpdateGranary{
		Id:        granaryId,
		UpdatedTs: currentTs,
	}

	rs := []int32{}
	updateResource := false

	for _, field := range request.UpdateMask.Paths {
		if field == "name" {
			update.Name = req.Name
		} else if field == "avatar" {
			update.Avatar = req.Avatar
		} else if field == "description" {
			update.Description = req.Description
		} else if field == "tags" {
			cb, e := json.Marshal(req.Tags)
			if e != nil {
				return nil, status.Errorf(codes.Internal, "failed to marshal tags: %v", e)
			}
			update.Tags = string(cb)
		} else if field == "sources" {
			if req.Sources == nil || len(req.Sources) == 0 {
				update.Sources = "{}"
			} else {

				for _, source := range req.Sources {
					rid, e := utils.ExtractIDFromUrn(utils.ResourceUrnPrefix, source)
					if e != nil {
						return nil, e
					}
					rs = append(rs, rid)
				}
				cb, e := json.Marshal(rs)
				if e != nil {
					return nil, e
				}
				update.Sources = string(cb)
			}
			updateResource = true
		} else if field == "options" {
			cb, e := json.Marshal(req.Options)
			if e != nil {
				return nil, status.Errorf(codes.Internal, "failed to marshal options: %v", e)
			}
			update.Options = string(cb)
		} else if field == "description" {
			update.Description = req.Description
		} else if field == "llm_model" {
			id, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, req.LlmModel)
			if err != nil {
				return nil, status.Errorf(codes.InvalidArgument, "invalid llm model: %s", req.LlmModel)
			}
			update.LlmModelId = id
		} else if field == "embedding_model" {
			id, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, req.EmbeddingModel)
			if err != nil {
				return nil, status.Errorf(codes.InvalidArgument, "invalid llm model: %s", req.LlmModel)
			}
			update.EmbeddingModelId = id
		} else {
			return nil, status.Errorf(codes.InvalidArgument, "invalid update path: %s", field)
		}
	}

	_, err = s.Gds.Handler.UpdateGranary(ctx, update)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to update granary: %v", err)
	}

	updated, err := s.Gds.Handler.GetGranary(ctx, granaryId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get granary: %v", err)
	}

	if updateResource {
		err = s.updateGranaryResources(ctx, granaryId, rs)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to update granary resources: %v", err)
		}
	}
	return convertGranaryFromDB(updated), nil
}

func (s *APIV1Service) DeleteGranary(ctx context.Context, request *v1pb.DeleteGranaryRequest) (*emptypb.Empty, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	id, err := utils.ExtractIDFromUrn(utils.GranaryUrnPrefix, request.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid urn: %v", err)
	}

	granary, err := s.Gds.Handler.GetGranary(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get granary, error: %+v", err)
	}

	if granary == nil {
		return nil, status.Errorf(codes.NotFound, "granary not found")
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() {
		if granary.TenantID != currentUser.TenantID {
			return nil, status.Errorf(codes.PermissionDenied, "permission denied")
		}

		if granary.CreatorID == currentUser.ID || currentUser.Role == cpb.Role_TENANT_ADMIN.String() {
			slog.Info("delete granary", "granary", granary.Name)
		} else {
			return nil, status.Errorf(codes.PermissionDenied, "permission denied")
		}
	}

	err = s.Gds.Handler.DeleteGranaryResource(ctx, &apb.FindGranaryResource{Gid: id})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete granary resource, error: %+v", err)
	}

	err = s.Gds.Handler.DeleteGranary(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to delete granary, error: %+v", err)
	}

	deleteRag := func() {
		rgb := ragBuildMap[request.Urn]
		if rgb == nil {
			p := &process.Process{Gds: s.Gds}
			ragConfig := p.NewRagConfig()
			rgb, err = p.NewRagBuilder(ragConfig)
			if err != nil {
				slog.Error("failed to create rag builder", slog.Any("error", err))
				return
			}
		}

		rgb.DeleteRag(granary.DocID)
	}

	go deleteRag()

	return nil, nil
}
func (s *APIV1Service) GetGranary(ctx context.Context, request *v1pb.GetGranaryRequest) (*v1pb.Granary, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	id, err := utils.ExtractIDFromUrn(utils.GranaryUrnPrefix, request.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid urn: %v", err)
	}

	granary, err := s.Gds.Handler.GetGranary(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get granary, error: %+v", err)
	}

	if granary == nil {
		return nil, status.Errorf(codes.NotFound, "granary not found")
	}

	if currentUser.TenantID > 0 && granary.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	return convertGranaryFromDB(granary), nil
}

var GranaryFilterCELAttributes = []cel.EnvOption{
	cel.Variable("name", cel.StringType),
	cel.Variable("sources", cel.ListType(cel.StringType)),
	cel.Variable("tags", cel.ListType(cel.StringType)),
}

type GranaryFilter struct {
	name    string
	sources []string
	tags    []string
}

func findGranayField(callExpr *expr.Expr_Call, filter *GranaryFilter) {
	if len(callExpr.Args) == 2 {
		idExpr := callExpr.Args[0].GetIdentExpr()
		if idExpr != nil {
			if idExpr.Name == "name" {
				filter.name = callExpr.Args[1].GetConstExpr().GetStringValue()
			}
			return
		} else if idExpr.Name == "tags" {
			for _, expr := range callExpr.Args[1].GetListExpr().GetElements() {
				value := expr.GetConstExpr().GetStringValue()
				filter.tags = append(filter.tags, value)
			}
			return
		} else if idExpr.Name == "sources" {
			for _, expr := range callExpr.Args[1].GetListExpr().GetElements() {
				value := expr.GetConstExpr().GetStringValue()
				filter.sources = append(filter.sources, value)
			}
			return
		}
	}
	for _, arg := range callExpr.Args {
		callExpr := arg.GetCallExpr()
		if callExpr != nil {
			findGranayField(callExpr, filter)
		}
	}
}

func parseGranaryFilter(expression string) (*GranaryFilter, error) {
	e, err := cel.NewEnv(GranaryFilterCELAttributes...)
	if err != nil {
		return nil, err
	}
	ast, issues := e.Compile(expression)
	if issues != nil {
		return nil, errors.Errorf("found issue %v", issues)
	}
	filter := &GranaryFilter{}
	expr, err := cel.AstToParsedExpr(ast)
	if err != nil {
		return nil, err
	}
	callExpr := expr.GetExpr().GetCallExpr()
	findGranayField(callExpr, filter)
	return filter, nil
}

func buildGranaryFindWithFilter(find *apb.FindGranary, filter string) error {
	if find == nil {
		find = &apb.FindGranary{}
	}

	if filter != "" {
		pfilter, err := parseGranaryFilter(filter)
		if err != nil {
			return status.Errorf(codes.InvalidArgument, "invalid filter: %v", err)
		}
		if len(pfilter.name) > 0 {
			find.Name = pfilter.name
		}

		if len(pfilter.sources) > 0 {
			find.Sources = pfilter.sources
		}

		if len(pfilter.tags) > 0 {
			find.Tags = pfilter.tags
		}
	}

	return nil
}

func (s *APIV1Service) ListGranaries(ctx context.Context, request *v1pb.ListGranariesRequest) (*v1pb.ListGranariesResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)

	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	find := &apb.FindGranary{}
	buildGranaryFindWithFilter(find, request.Filter)

	var limit, offset int32
	if request.PageToken != nil {
		limit = request.PageToken.Limit
		offset = request.PageToken.Offset
	} else {
		limit = request.PageSize
	}
	if limit <= 0 {
		limit = DefaultPageSize
	}

	if currentUser.Role != cpb.Role_HOST.String() && currentUser.Role != cpb.Role_ADMIN.String() {
		find.TenantId = currentUser.TenantID
	}

	limitPlusOne := limit
	find.Limit = limitPlusOne
	find.Offset = offset

	granaries, count, err := s.Gds.Handler.ListGranaries(ctx, find)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list granary: %v", err)
	}
	granaryMessages := []*v1pb.Granary{}

	if int32(len(granaries)) == limitPlusOne {
		granaries = granaries[:limit]
	}

	userMap := make(map[int32]*v1pb.User)
	for _, g := range granaries {
		pMsg := convertGranaryFromDB(g)
		user := userMap[g.CreatorID]
		if user == nil {
			u, _ := s.Gds.Handler.GetUser(ctx, &apb.FindUser{Id: g.CreatorID})
			if u == nil {
				continue
			}
			user = s.convertUserFromStore(u)

			if err != nil {
				return nil, errors.Wrap(err, "failed to convert user")
			}
		}

		pMsg.Creator = user

		userMap[g.CreatorID] = user
		granaryMessages = append(granaryMessages, pMsg)
	}

	response := &v1pb.ListGranariesResponse{
		TotalSize: count,
		Granaries: granaryMessages,
		NextPageToken: &cpb.PageToken{
			Limit:  int32(limit),
			Offset: int32(offset + limit),
		},
	}

	return response, nil
}

func (s *APIV1Service) ListGranariesTopics(ctx context.Context, request *v1pb.ListGranariesTopicsRequest) (*v1pb.ListGranariesTopicsResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)

	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	granaryId, err := utils.ExtractIDFromUrn(utils.GranaryUrnPrefix, request.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid user name: %v", err)
	}

	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	granary, err := s.Gds.Handler.GetGranary(ctx, granaryId)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	if granary == nil {
		return nil, status.Errorf(codes.NotFound, "user not found")
	}

	if currentUser.TenantID > 0 && currentUser.TenantID != granary.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	if granary.Status != v1pb.GranaryStatus_ENDED.String() {
		return nil, status.Errorf(codes.PermissionDenied, "granary is not active")
	}
	// p := &process.Process{Gds: s.Gds}
	// ragConfig := p.NewRagConfig()
	// rgb, err := p.NewRagBuilder(ragConfig)
	// if err != nil {
	// 	return nil, status.Errorf(codes.Internal, "failed to create rag builder: %v", err)
	// }

	// summarites, err := rgb.QuerySummarTopic(granary.DocID)
	// if err != nil {
	// 	return nil, status.Errorf(codes.Internal, "failed to query summar topic: %v", err)
	// }

	// topics := []*cpb.TopicDef{}
	// for _, summarite := range summarites {
	// 	topics = append(topics, summarite.Summary)
	// }

	topics, count, err := s.FetchTopic(granary)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to fetch topic: %v", err)
	}

	topicResponce := &v1pb.ListGranariesTopicsResponse{
		TotalSize: count,
		Topics:    topics,
	}
	return topicResponce, nil

}

func (s *APIV1Service) FetchTopic(granary *model.Granary) ([]*cpb.TopicDef, int32, error) {
	p := &process.Process{Gds: s.Gds}
	ragConfig := p.NewRagConfig()
	rgb, err := p.NewRagBuilder(ragConfig)
	if err != nil {
		return nil, 0, status.Errorf(codes.Internal, "failed to create rag builder: %v", err)
	}

	summarites, err := rgb.QuerySummarTopic(granary.DocID)
	if err != nil {
		return nil, 0, status.Errorf(codes.Internal, "failed to query summar topic: %v", err)
	}

	topics := []*cpb.TopicDef{}
	for _, summarite := range summarites {
		topics = append(topics, summarite.Summary)
	}
	return topics, int32(len(summarites)), nil
}

func (s *APIV1Service) CtrlGranary(ctx context.Context, request *v1pb.CtrlGranaryRequest) (*v1pb.CtrlGranaryResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)

	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	id, err := utils.ExtractIDFromUrn(utils.GranaryUrnPrefix, request.Urn)

	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid urn: %v", err)
	}

	granary, err := s.Gds.Handler.GetGranary(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get granary: %v", err)
	}
	if granary.CreatorID != currentUser.ID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	startGranary := func(input []int32) {
		slog.Info("start granary", slog.Any("granary input", input))
		if len(input) == 0 {
			slog.Info("granary sources is empty")

			return
		}
		nctx, _ := UserCtxClone(ctx)
		err := s.startGranary(nctx, request, input, granary)
		if err != nil {
			slog.Error("failed to start granary:", slog.Any("err", err))
			granary.Status = v1pb.GranaryStatus_FAILED.String()
			s.Gds.Handler.UpdateGranarySimple(nctx, granary)
		}
	}
	switch request.Cmd {
	case v1pb.GranaryCmd_START:
		var inputInfo []int32
		e := json.Unmarshal([]byte(granary.Sources), &inputInfo)
		if e != nil {
			return nil, status.Errorf(codes.InvalidArgument, "invalid sources: %v", e)
		}
		go startGranary(inputInfo)
	case v1pb.GranaryCmd_INCREMENTAL:
		var inputInfo []int32
		grs, e := s.Gds.Handler.ListGranaryResources(ctx, &apb.FindGranaryResource{Gid: granary.ID})
		if e != nil {
			return nil, status.Errorf(codes.Internal, "failed to list granary resources: %v", e)
		}

		for _, gr := range grs {
			if gr.Status != v1pb.GranaryStatus_ENDED.String() {
				inputInfo = append(inputInfo, gr.ResourceID)
			}
		}
		go startGranary(inputInfo)
	case v1pb.GranaryCmd_STOP:
		return s.stopGranary(ctx, granary)
	default:
	}

	granary.Status = v1pb.GranaryStatus_STARTING.String()
	s.Gds.Handler.UpdateGranarySimple(ctx, granary)
	e := status.Errorf(codes.OK, "start process documents")
	return buildGranaryCtrlResponse(codes.OK, e)
}

func (s *APIV1Service) startGranary(ctx context.Context, request *v1pb.CtrlGranaryRequest, inputInfo []int32, granary *model.Granary) error {
	update := false
	if request.LlmModel != "" {
		lmid, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, request.LlmModel)
		if err != nil {
			return status.Errorf(codes.InvalidArgument, "invalid llm model urn: %v", err)
		}
		granary.LlmModelID = lmid
		update = true
	}

	if request.EmbeddingModel != "" {
		emid, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, request.EmbeddingModel)
		if err != nil {
			return status.Errorf(codes.InvalidArgument, "invalid embedding model urn: %v", err)

		}
		granary.EmbeddingModelID = emid
		update = true
	}
	if update {
		err := s.Gds.Handler.UpdateGranarySimple(ctx, granary)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to update granary: %v", err)
		}
	}

	if granary.LlmModelID < 1 || granary.EmbeddingModelID < 1 {
		return status.Errorf(codes.InvalidArgument, "llm model or embedding model is not set")
	}

	p := &process.Process{Gds: s.Gds}
	ragConfig := p.NewRagConfig()
	rgb, e := p.SetRagAssistant(granary.LlmModelID, granary.EmbeddingModelID, ragConfig)
	if e != nil {
		return status.Errorf(codes.Internal, "failed to set rag assistant: %v", e)
	}
	rgb.GranaryId = granary.ID

	ragBuildMap[request.Urn] = rgb

	increment := true
	if request.Cmd == v1pb.GranaryCmd_START {
		increment = false
	}

	return s.startProcessDocuments(ctx, inputInfo, granary, rgb, increment)
}

func (s *APIV1Service) stopGranary(ctx context.Context, granary *model.Granary) (*v1pb.CtrlGranaryResponse, error) {
	granaryUrn := utils.BuildUrn(utils.GranaryUrnPrefix, granary.ID)
	rgb := ragBuildMap[granaryUrn]

	if rgb != nil {
		rgb.Cancel()
		rgb = nil
	}

	if granary.Status == "STOPPED" || granary.Status == "ENDED" || granary.Status == "FAILED" {
		return buildGranaryCtrlResponse(codes.OK, status.Errorf(codes.OK, "granary already stopped"))
	}

	granary.Status = "STOPPED"

	e := s.Gds.Handler.UpdateGranarySimple(ctx, granary)
	if e != nil {
		return nil, status.Errorf(codes.Internal, "failed to update granary: %v", e)
	}

	e = s.Gds.Handler.StopGranaryResource(ctx, &apb.FindGranaryResource{Gid: granary.ID})
	if e != nil {
		return nil, status.Errorf(codes.Internal, "failed to stop granary resources: %v", e)
	}
	return buildGranaryCtrlResponse(codes.OK, status.Errorf(codes.OK, "granary stop success"))
}

func (s *APIV1Service) startProcessDocuments(ctx context.Context, fileUrns []int32, granary *model.Granary, rgb *rag.RagBuilder, increment bool) error {
	docs := []*apb.RagResource{}
	for _, urn := range fileUrns {
		path, fileType, err := s.GetResourcePath(ctx, urn)
		slog.Info("get resource path", slog.Any("path", path), slog.Any("fileType", fileType), slog.Any("urn", urn))
		if err != nil {
			slog.Error("failed to get resource path", slog.Any("error", err))
			continue
		}
		d := &apb.RagResource{
			RsType: fileType,
			RsPath: path,
			RsId:   urn,
			Gid:    granary.ID,
		}
		docs = append(docs, d)
	}

	if len(docs) == 0 {
		granary.Status = v1pb.GranaryStatus_FAILED.String()
		s.Gds.Handler.UpdateGranarySimple(ctx, granary)
		return status.Errorf(codes.InvalidArgument, "granary sources is empty")
	}

	var options map[string]string

	if granary.Options != "" && granary.Options != "{}" {
		e := json.Unmarshal([]byte(granary.Options), &options)
		if e != nil {
			slog.Error("failed to unmarshal options", slog.Any("error", e))
			return e
		}
	}

	if increment {
		options["drop"] = "false"
	}

	go rgb.ProcessDoc(docs, granary.DocID, options, s.updateGranaryStatus, s.updateGranaryResourceStatus)
	return nil
}

func (s *APIV1Service) RagTest(ctx context.Context, request *v1pb.RagTestRequest) (*v1pb.RagTestResponse, error) {
	id, e := utils.ExtractIDFromUrn(utils.GranaryUrnPrefix, request.Urn)
	if e != nil {
		return nil, status.Errorf(codes.Internal, "failed to extract id from urn: %v", e)
	}

	granary, e := s.Gds.Handler.GetGranary(ctx, id)
	if e != nil {
		return nil, status.Errorf(codes.Internal, "failed to get granary: %v", e)
	}

	res := &v1pb.RagTestResponse{
		Tid: shortuuid.New(),
	}

	if request.InteractType == cpb.InteractType_CHAT {
		s.registerGranaryWebsocket(ctx, granary, request.RagMode, request.Content, request.LlmModel, res.Tid)
	} else if request.InteractType == cpb.InteractType_QUERY {
		return s.ProcessCompleteRagTest(ctx, granary, request.RagMode, request.Content, request.LlmModel)
	}

	return res, nil
}

func (s *APIV1Service) ProcessCompleteRagTest(ctx context.Context, srcGranary *model.Granary, ragMode cpb.RagMode, content string, modelUrn string) (*v1pb.RagTestResponse, error) {
	p := process.NewProcess(s.Gds)
	ragClientConfig := p.NewRagConfig()
	llmModelId := srcGranary.LlmModelID
	if modelUrn != "" {
		i, err := utils.ExtractIDFromUrn(utils.ModelUrnPrefix, modelUrn)
		if err == nil {
			llmModelId = i
		}
	}

	rgb, e := p.SetRagAssistant(llmModelId, srcGranary.EmbeddingModelID, ragClientConfig)
	if e != nil {
		slog.Error("failed to set rag assistant", slog.Any("err", e))
		return nil, status.Errorf(codes.Internal, "failed to set rag assistant: %v", e)
	}
	var docs []string
	docs = append(docs, srcGranary.DocID)

	prompt, err := rgb.QueryLLMRagPrompt(content, docs, true, ragMode)
	if err != nil {
		slog.Error("failed to get rag prompt", slog.Any("err", err))
		return nil, status.Errorf(codes.Internal, "failed to get rag prompt: %v", err)
	}

	messages := []llm.RoleContent{}
	messages = append(messages, llm.RoleContent{
		Role:    llm.RoleSys,
		Content: prompt,
	})

	messages = append(messages, llm.RoleContent{
		Role:    llm.RoleUser,
		Content: content,
	})

	c := factory.NewChatClient(ragClientConfig.ChatConfig)

	req := &llm.CompleteReq{Msgs: messages}
	answer, err := c.Complete(ctx, req)
	if err != nil {
		slog.Error("failed to query rag assistant", slog.Any("err", err))
		return nil, status.Errorf(codes.Internal, "failed to query rag assistant: %v", err)
	}

	if answer == nil {
		return nil, status.Errorf(codes.Internal, "answer is nil")
	}

	res := &v1pb.RagTestResponse{
		Content: answer.GetContent(),
		Tid:     shortuuid.New(),
	}
	return res, nil
}

type RagService struct {
	userId int32
	api    *APIV1Service
	ctx    context.Context
	rgb    *rag.RagBuilder
	rgc    *rag.RagClientConfig
	docs   []string
}

func (s *RagService) BuildChatMessage(phase v1pb.Phase, content string, askId string) *v1pb.Exchange {
	return &v1pb.Exchange{Phase: phase, Content: content, AskId: askId}
}

func (s *RagService) sendOutMessage(ctx context.Context, outMessageChan chan<- *v1pb.Exchange, message *v1pb.Exchange) bool {
	select {
	case <-ctx.Done():
		slog.Warn("chat service context cancelled, stopping message send")
		return false
	default:
		// 如果 context 没有取消，尝试发送消息
		select {
		case outMessageChan <- message:
			return true
		default:
			slog.Warn("channel full or closed")
			return false
		}
	}
}

func (s *RagService) doChat(ragMode cpb.RagMode, content string, outMessageChan chan<- *v1pb.Exchange) {
	messages := []llm.RoleContent{}

	if s.rgb == nil {
		return
	}

	useDefault := true
	prompt, err := s.rgb.QueryLLMRagPrompt(content, s.docs, useDefault, ragMode)
	if err != nil {
		return
	}
	messages = append(messages, llm.RoleContent{
		Role:    llm.RoleSys,
		Content: prompt,
	})

	messages = append(messages, llm.RoleContent{
		Role:    llm.RoleUser,
		Content: content,
	})

	chatReq := &llm.ChatReq{Msgs: messages}
	chatReq.RspChan = make(chan string)
	chatReq.CompleteChan = make(chan bool)
	c := factory.NewChatClient(s.rgc.ChatConfig)
	go func() {
		select {
		case <-s.ctx.Done():
			slog.Info("recv cancel, return now.")
			return
		default:
			s.sendOutMessage(s.ctx, outMessageChan, s.BuildChatMessage(v1pb.Phase_Begin, ".", ""))
			c.Chat(s.ctx, chatReq)
		}
	}()
	var assistantmsgs []string
	for {
		select {
		case <-s.ctx.Done():
			slog.Info("recv cancel. break now.")
			return
		case rsp := <-chatReq.RspChan:
			if s.sendOutMessage(s.ctx, outMessageChan, s.BuildChatMessage(v1pb.Phase_Continue, rsp, "")) {
				break
			}
		case complete := <-chatReq.CompleteChan:
			if complete {
				messages = append(messages, llm.RoleContent{
					Role:    llm.RoleAssistant,
					Content: strings.Join(assistantmsgs, ""),
				})

				assistantmsgs = nil
				if !s.sendOutMessage(s.ctx, outMessageChan, s.BuildChatMessage(v1pb.Phase_End, "", "")) {
					return
				}

				if len(messages) > 110 {
					if !s.sendOutMessage(s.ctx, outMessageChan, s.BuildChatMessage(v1pb.Phase_Close, "记录超过50条，请新建会话", "askid")) {
						return
					}
				}
				return
			}
		case <-time.After(5 * time.Second):
			// 超时处理
		}
	}

}

func (s *RagService) processRagRequest(ctx context.Context, srcGranary *model.Granary, ragMode cpb.RagMode, modelUrn string, inMessageChan <-chan *v1pb.Exchange, outMessageChan chan<- *v1pb.Exchange) {
	p := process.NewProcess(s.api.Gds)
	ragClientConfig := p.NewRagConfig()
	rgb, e := p.SetRagAssistant(srcGranary.LlmModelID, srcGranary.EmbeddingModelID, ragClientConfig)
	s.rgb = rgb
	s.rgc = ragClientConfig
	if e != nil {
		slog.Error("failed to set rag assistant", slog.Any("err", e))
		return
	}
	var docs []string
	docs = append(docs, srcGranary.DocID)
	cctx, cancel := context.WithCancel(context.Background())
	s.ctx = cctx
	for {
		select {
		case <-ctx.Done():
			slog.Info("rag service for chat service is done")
			cancel()
			return
		case inMsg, ok := <-inMessageChan:
			if !ok {
				slog.Info("recv error.")
				return
			}

			if inMsg.Phase == v1pb.Phase_End {
				slog.Info("Received end message:", slog.Any("id", inMsg.AskId))
				cancel()
				cctx, cancel = context.WithCancel(context.Background())
				break
			}

			if inMsg.Content == "" {
				continue
			}

			msgcontent := s.BuildChatMessage(v1pb.Phase_Prepare, "正在分析中...", "")
			s.sendOutMessage(cctx, outMessageChan, msgcontent)

			if inMsg.Mode != cpb.RagMode_RAG_MODE_UNSPECIFIED {
				ragMode = inMsg.Mode
			}

			go func() {
				s.doChat(ragMode, inMsg.Content, outMessageChan)
			}()
		}

	}
}

func (s *APIV1Service) handleRagWebSocketMessages(ctx context.Context, cancel context.CancelFunc, gid string, conn *websocket.Conn, granary *model.Granary, ragMode cpb.RagMode, inMessageChan chan *v1pb.Exchange) {
	// 添加互斥锁来保护WebSocket写操作
	writeMutex := s.WsManager.MutexMap[gid]
	sendInputMessage := func(askId string, content string, phase v1pb.Phase, mode cpb.RagMode) {
		inMsg := &v1pb.Exchange{
			AskId:   askId,
			Content: content,
			Phase:   phase,
			Mode:    mode,
		}
		select {
		case inMessageChan <- inMsg:
			slog.Info("Message sent to inMessageChan", "message", content)
		case <-ctx.Done():
			return
		default:
			slog.Warn("Failed to send message to inMessageChan: channel full or closed")
		}
	}
	for {
		select {
		case <-ctx.Done():
			return
		default:
			_, msg, err := conn.ReadMessage()
			if err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					slog.Warn("WebSocket close error:", slog.Any("err", err))
					cancel()
				} else {
					slog.Warn("WebSocket other error:", slog.Any("err", err))
					cancel()
				}
				return
			}

			if string(msg) == "xxpingxx" {
				writeMutex.Lock()
				err := conn.WriteMessage(websocket.TextMessage, []byte("xxpongxx"))
				writeMutex.Unlock()
				if err != nil {
					slog.Warn("Error sending pong:", slog.Any("err", err))
					continue
				}
				continue
			}

			var chatMessage v1pb.Exchange
			if err := proto.Unmarshal(msg, &chatMessage); err != nil {
				slog.Warn("Error unmarshaling message:", slog.Any("err", err))
				continue
			}
			ragMode := ragMode
			if chatMessage.Mode != cpb.RagMode_RAG_MODE_UNSPECIFIED {
				ragMode = chatMessage.Mode
			}
			if chatMessage.Phase == v1pb.Phase_HELLO {
				content := chatMessage.Content
				if content == "" {
					content = s.WsManager.InitMessage[chatMessage.AskId]
					if content == "" {
						slog.Info("hello phase message is empty ", chatMessage.AskId, content)
						continue
					}
					slog.Info("hello init granary message", chatMessage.AskId, content)
					sendInputMessage(chatMessage.AskId, content, chatMessage.Phase, ragMode)
					continue
				}

			} else if chatMessage.Phase == v1pb.Phase_End {
				sendInputMessage(chatMessage.AskId, "", chatMessage.Phase, ragMode)
				continue
			} else if chatMessage.Phase == v1pb.Phase_Close {
				slog.Info("close phase message")
				cancel()
				return
			}

			trimmedContent := strings.TrimSpace(chatMessage.Content)

			if trimmedContent == "" {
				continue
			}

			sendInputMessage(shortuuid.New(), chatMessage.Content, chatMessage.Phase, ragMode)
		}
	}
}

func (s *APIV1Service) registerGranaryWebsocket(ctx context.Context, srcGranary *model.Granary, ragMode cpb.RagMode, content string, modelUrn string, gid string) {
	s.WsManager.InitMessage[gid] = content

	s.WsManager.Register(gid, func(conn *websocket.Conn) {
		slog.Info("WebSocket connection established for granary:", slog.Any("UID", gid))
		outMessageChan := make(chan *v1pb.Exchange, 5) // 缓冲大小可以根据需要调整
		inMessageChan := make(chan *v1pb.Exchange, 5)  // 缓冲大小可以根据需要调整

		nctx, cancel := context.WithCancel(context.Background())
		nctx = context.WithValue(nctx, userIdContextKey, srcGranary.CreatorID)
		userId, err := s.GetCurrentUserId(ctx) // 获取当前用户ID
		if err != nil {
			slog.Error("Failed to get current user ID", slog.Any("error", err))
			return
		}

		ragService := &RagService{
			// done:      make(chan struct{}),
			userId: userId,
			api:    s,
			docs:   []string{srcGranary.DocID},
		}

		defer conn.Close()
		defer close(inMessageChan)  // 确保在连接关闭时关闭输入通道
		defer close(outMessageChan) // 确保在处理完成后关闭 channel
		defer cancel()              // 确保在处理完成后取消 context
		defer s.WsManager.UnRegister(gid)

		go ragService.processRagRequest(nctx, srcGranary, ragMode, modelUrn, inMessageChan, outMessageChan)
		go s.handleCommonWebSocketMessages(nctx, cancel, gid, conn, ragMode, inMessageChan)
		s.handleExchangeMessage(nctx, cancel, gid, conn, outMessageChan)
	})
}

func buildGranaryCtrlResponse(code codes.Code, err error) (*v1pb.CtrlGranaryResponse, error) {
	errMsg := "Success"
	if err != nil {
		errMsg = err.Error()
	}

	rsp := &v1pb.CtrlGranaryResponse{
		Code:    int32(code),
		Message: errMsg,
	}

	return rsp, err
}

func (s *APIV1Service) updateGranaryStatus(ctx context.Context, id int32, status v1pb.GranaryStatus, summary string, tags []string) error {
	granary, err := s.Gds.Handler.GetGranary(ctx, id)
	if err != nil {
		return err
	}
	granary.Status = status.String()
	if status == v1pb.GranaryStatus_ENDED {
		granary.Summary = summary
		b, e := json.Marshal(tags)
		if e != nil {
			return e
		}
		granary.Tags = string(b)
	}
	err = s.Gds.Handler.UpdateGranarySimple(ctx, granary)
	if err != nil {
		return err
	}
	if status == v1pb.GranaryStatus_ENDED || status == v1pb.GranaryStatus_STOPPED {
		delete(ragBuildMap, utils.BuildUrn(utils.GranaryUrnPrefix, granary.ID))
	}
	return nil
}

func (s *APIV1Service) updateGranaryResourceStatus(ctx context.Context, gid int32, uid int32, status v1pb.GranaryStatus, summary string, tags []string) error {
	gr := &model.GranaryResource{
		GranaryID:  gid,
		ResourceID: uid,
		Status:     status.String(),
		Summary:    summary,
	}

	if tags != nil && len(tags) > 0 {
		b, e := json.Marshal(tags)
		if e != nil {
			return e
		}
		gr.Tags = string(b)
	}
	err := s.Gds.Handler.UpsertGranaryResource(ctx, gr)
	if err != nil {
		return err
	}
	return nil
}

func (s *APIV1Service) ListGranaryResources(ctx context.Context, request *v1pb.ListGranaryResourcesRequest) (*v1pb.ListGranaryResourcesResponse, error) {
	currentUser, err := s.GetCurrentUser(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get user: %v", err)
	}

	id, err := utils.ExtractIDFromUrn(utils.GranaryUrnPrefix, request.Urn)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid urn: %v", err)
	}

	granary, err := s.Gds.Handler.GetGranary(ctx, id)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to get granary, error: %+v", err)
	}

	if granary == nil {
		return nil, status.Errorf(codes.NotFound, "granary not found")
	}

	if currentUser.TenantID > 0 && granary.TenantID != currentUser.TenantID {
		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
	}

	grs, err := s.Gds.Handler.ListGranaryResources(ctx, &apb.FindGranaryResource{Gid: granary.ID})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list granary resources, error: %+v", err)
	}

	grStates := make([]*v1pb.GranaryResourcecState, len(grs))
	for _, gr := range grs {
		state := &v1pb.GranaryResourcecState{
			DocUrn:  utils.BuildUrn(utils.ResourceUrnPrefix, gr.ResourceID),
			Status:  v1pb.GranaryStatus(v1pb.GranaryStatus_value[gr.Status]),
			Summary: gr.Summary,
		}

		r, e := s.Gds.Handler.GetResource(ctx, &apb.FindResource{Id: gr.ResourceID})
		if e != nil {
			return nil, status.Errorf(codes.Internal, "failed to get resource, error: %+v", e)
		}

		state.Resource = s.convertResourceFromDB(ctx, r)

		if gr.Tags != "" {
			var tags []string
			e := json.Unmarshal([]byte(gr.Tags), &tags)
			if e == nil {
				state.Tags = tags
			}
		}
		grStates = append(grStates, state)
	}

	res := &v1pb.ListGranaryResourcesResponse{
		TotalSize: int32(len(grStates)),
		Resources: grStates,
	}
	return res, nil
}
